import os
import argparse
import json
from openai import OpenAI

def optimize_prompt_with_api(instruction, context=None):
    """
    Uses the OpenAI API to optimize a given prompt.

    Args:
        instruction (str): The main instruction to be optimized.
        context (str, optional): Additional context for the instruction. Defaults to None.

    Returns:
        dict: The optimized prompt in a JSON object, or an error message.
    """
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        return {"error": "OPENAI_API_KEY environment variable not set."}

    client = OpenAI(api_key=api_key)

    system_prompt = "你是一个提示词优化专家，旨在将用户输入优化得更清晰、更有效、更具结构化。请以JSON格式返回优化结果，返回格式必须是：{\"optimized_prompt\": \"优化后的提示词内容\"}"

    if context:
        user_prompt = f"请优化以下指令：\n\n### 指令:\n{instruction}\n\n### 上下文:\n{context}"
    else:
        user_prompt = f"请优化以下指令：\n\n### 指令:\n{instruction}"

    try:
        print("向 OpenAI API 发送请求...")
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            # 强制模型返回JSON对象，这是最关键的功能
            response_format={"type": "json_object"},
            temperature=0.7,
        )
        
        print("成功接收到 API 响应。")
        # API保证了返回的是一个有效的JSON字符串
        json_response = json.loads(response.choices[0].message.content)
        return json_response

    except Exception as e:
        return {"error": f"An API error occurred: {e}"}

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Optimize a prompt using the OpenAI GPT-4o API.")
    parser.add_argument("--instruction", type=str, required=True, help="The instruction to be optimized.")
    parser.add_argument("--context", type=str, default=None, help="Optional context for the instruction.")
    
    args = parser.parse_args()
    
    optimized_result = optimize_prompt_with_api(args.instruction, args.context)
    
    # 使用json.dumps来格式化输出，确保中文正常显示
    print("\n--- 优化结果 ---\n")
    print(json.dumps(optimized_result, indent=2, ensure_ascii=False)) 